#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <float.h>
#include <cuda.h>
#include "cublas.h"
#include <time.h>
#include <assert.h>

#define ceil(x)	( x - (x % 1) + 1 )

#define INF 		65535
#define NUM_PWARPS	16
//#define NUM_VWARPS	(NUM_PWARPS * (32/WARP_SIZE) )
#define NUM_VWARPS	(NUM_PWARPS * (32/32) )
#define MIN(X,Y) (((X)<(Y))?(X):(Y))
#define MAX(X,Y) (((X)>(Y))?(X):(Y))

typedef struct graphstruct { // A graph in compressed-adjacency-list (CSR) form
    int nv;            // number of vertices
    int ne;            // number of edges
    int *nbr;          // array of neighbors of all vertices
    int *firstnbr;     // index in nbr[] of first neighbor of each vtx
} graph;

template<int WARP_SIZE>
struct warp_mem 
{
	int levels[WARP_SIZE];
	int nodes[WARP_SIZE+1];
	int scratch ;
};

template<int WARP_SIZE> __global__ void warp_baseline_kn(int, bool *, int *, int , int , int *, int *, char* );
//template<int> __global__ void warp_bfs_kernel(int *,int *, int *, int , bool *,int *);
template<int> __device__ void memcpy_SIMD(int , int , int* , int* );
template<int> __device__ void expand_bfs_SIMD(int , int , int* , int *, int , bool *, int *);

//--------------------------------------------------------------
// For WARP-based SIMD operations
//--------------------------------------------------------------
template <typename T, int WARP_SIZE>
__device__
inline static void gpu_memcpy_SIMD(T* d_dst, T* d_src, int cnt, int warp_offset)
{
	// each thread copies one element
	for(int i=warp_offset; i < cnt; i+= WARP_SIZE)
	{
		d_dst[i] = d_src[i];
	}
    __threadfence_block();
}

template <typename T2, typename T, int WARP_SIZE>
__device__ static inline void gpu_memcpy_SIMD2(T2* d_dst, T* d_src, int cnt, int warp_offset)
{
	// each thread copies one element
	for(int i=warp_offset; i < cnt; i+= WARP_SIZE)
	{
		d_dst[i] = (T2) (d_src[i]);
	}
    __threadfence_block();
}

template<int WARP_SIZE>
__device__ static bool expand_frontier_gpu_SIMD(int M, int from, int to, int level, int * d_edges, int* d_level, int warp_offset, bool* d_finished)
{
    bool finished = true;
	for(int k = from + warp_offset; k < to; k += (WARP_SIZE)) {
		int nbr = d_edges[k]; // this read is coalesced
		if (d_level[nbr] == INF) {	// only this final read-write is scattererd
			d_level[nbr] = level + 1;
#if USE_OLD_GPU_METHOD
            *d_finished = false;
#endif
            finished = false;
		}
	}   
    __threadfence_block();
    return finished;
}


int read_edge_list (int **tailp, int **headp) {
    int max_edges = 100000000;
    int nedges, nr, t, h;
    *tailp = (int *) calloc(max_edges, sizeof(int));
    *headp = (int *) calloc(max_edges, sizeof(int));
    nedges = 0;
    nr = scanf("%i %i",&t,&h);
    while (nr == 2) {
        if (nedges >= max_edges) {
            printf("Limit of %d edges exceeded.\n",max_edges);
            exit(1);
        }
        (*tailp)[nedges] = t;
        (*headp)[nedges++] = h;
        nr = scanf("%i %i",&t,&h);
    }
    return nedges;
}

graph * graph_from_edge_list (int *tail, int* head, int nedges) {
    graph *G;
    int i, e, v, maxv;
    G = (graph *) calloc(1, sizeof(graph));
    G->ne = nedges;
    maxv = 0;

    // count vertices
    for (e = 0; e < G->ne; e++) {
        if (tail[e] > maxv) maxv = tail[e];
        if (head[e] > maxv) maxv = head[e];
    }
    G->nv = maxv+1;
    G->nbr = (int *) calloc(G->ne, sizeof(int));
    G->firstnbr = (int *) calloc((G->nv)+1, sizeof(int));
    
    // count neighbors of vertex v in firstnbr[v+1],
    for (e = 0; e < G->ne; e++) G->firstnbr[tail[e]+1]++;
    
    // cumulative sum of neighbors gives firstnbr[] values
    for (v = 0; v < G->nv; v++) G->firstnbr[v+1] += G->firstnbr[v];
    
    // pass through edges, slotting each one into the CSR structure
    for (e = 0; e < G->ne; e++) {
        i = G->firstnbr[tail[e]]++;
        G->nbr[i] = head[e];
    }
    // the loop above shifted firstnbr[] left; shift it back right
    for (v = G->nv; v > 0; v--) G->firstnbr[v] = G->firstnbr[v-1];
    G->firstnbr[0] = 0;
    return G;
}


void print_CSR_graph (graph *G) {
    int vlimit = 20;
    int elimit = 50;
    int e,v;
    printf("\nGraph has %d vertices and %d edges.\n",G->nv,G->ne);
    printf("firstnbr =");
    if (G->nv < vlimit) vlimit = G->nv;
    for (v = 0; v <= vlimit; v++) printf(" %d",G->firstnbr[v]);
    if (G->nv > vlimit) printf(" ...");
    printf("\n");
    printf("nbr =");
    if (G->ne < elimit) elimit = G->ne;
    for (e = 0; e < elimit; e++) printf(" %d",G->nbr[e]);
    if (G->ne > elimit) printf(" ...");
    printf("\n\n");
}


int main(int argc, char **argv)
{
	graph *G;
	int *tail, *head ;
	int nedges , nlevels, startvtx ;
	cudaError_t err ;

	int *testvar_gpu ;
	char *d_debug ;

	//My variables
	int level ,i, *ctr; 
	bool *h_finished , *d_finished;
	int *d_nodes, *d_edges, *d_level , *h_level ; 

	h_finished = (bool *) malloc ( 1 * sizeof(bool));

	nedges = read_edge_list(&tail, &head);
	G = graph_from_edge_list(tail, head, nedges);
	
	h_level = (int *)malloc( (G->nv) * sizeof(int));

	for(i = 1 ; i < G->nv ; i++)
		h_level[i] = INF ;

	free(tail);
	free(head);
	print_CSR_graph(G);

	cudaMalloc( (void **)&d_finished, 1 * sizeof(bool)) ; assert(d_finished != NULL);
	cudaMalloc( (void **)&d_debug, 1 * sizeof(char)) ; assert(d_debug != NULL); 
	cudaMalloc( (void **)&testvar_gpu, 1 * sizeof(int)) ; assert(testvar_gpu != NULL);

	level = 0 ;

	err = cudaMalloc( (void **)&d_nodes,(G->nv+1) * sizeof(int)); assert(d_nodes!=NULL);
	err = cudaMalloc( (void **)&d_edges, (G->ne) * sizeof(int)); assert(d_edges!=NULL);
	err = cudaMalloc( (void **)&d_level, (G->nv) * sizeof(int)); assert(d_level!=NULL);

	int num_warps = NUM_VWARPS;	// 512 threads in TB
	int numthreadsperblock = num_warps * 32 ;
	int numblocks = MIN(16384, (G->nv + numthreadsperblock-1)/ numthreadsperblock); // number of threads blocks 
	numblocks = MAX(numblocks, 1);

	printf("Copying Graph from Host to Device :\n");

	err = cudaMemcpy(d_nodes,G->firstnbr, (G->nv+1) * sizeof(int) , cudaMemcpyHostToDevice); assert( err == cudaSuccess);
	err = cudaMemcpy(d_edges,G->nbr, (G->ne) * sizeof(int) ,  cudaMemcpyHostToDevice); assert( err == cudaSuccess);
	err = cudaMemcpy(d_level,h_level, (G->nv) * sizeof(int) , cudaMemcpyHostToDevice); assert( err == cudaSuccess);

	printf("GPU performing BFS ...\n");

	clock_t start=clock();

	do
	{
	    *h_finished = true ;

	    cudaMemcpy(d_finished,h_finished,1 * sizeof(bool), cudaMemcpyHostToDevice);
	    assert ( err == cudaSuccess);
	     warp_baseline_kn<32> <<< numblocks, numthreadsperblock >>> (
			level++, d_finished, d_level, G->nv, G->ne, d_edges, d_nodes, d_debug);
	    cudaMemcpy(h_finished,d_finished,1 * sizeof(bool), cudaMemcpyDeviceToHost);

	}
	while(!(*h_finished));
 
	printf("Copying result to host ...\n");
	cudaMemcpy(h_level , d_level, (G->nv) * sizeof(int) , cudaMemcpyDeviceToHost);

	printf("Time elapsed: %f\n", ((double)clock() - start) / CLOCKS_PER_SEC);
	printf("Number of levels = %d\n",level);

	ctr = (int *)malloc(level * sizeof(int));

	for(i = 0 ; i < level ; i++)
		ctr[i] = 0 ; 

	for(i = 0 ; i < (G->nv) ; i++)
	    if(h_level[i] < INF)
		ctr[h_level[i]]++;

	for(i = 0 ; i < level ; i++)
		printf("level %d vertices: %d\n",i, ctr[i]);


		printf("Levels array: \n");
	//for(i = 0 ; i < G->nv ; i++)
	//	printf("%d\t",h_level[i]);
	printf("\n");


	return 0 ; 
}

template<int WARP_SIZE>
__global__ void warp_baseline_kn(int level, bool* d_finished,  int* d_level, int N, int M, int* d_edges, int* d_nodes, char* d_debug)
{
	__shared__ warp_mem<WARP_SIZE> local_mem[NUM_VWARPS];

	int i = blockIdx.x * blockDim.x + threadIdx.x;
	const unsigned int grid_size = gridDim.x * blockDim.x;

	int warp_id = i / (WARP_SIZE);
	int warp_offset = threadIdx.x % (WARP_SIZE);
	int local_warp_id = threadIdx.x / WARP_SIZE;
	warp_mem<WARP_SIZE>* local = &(local_mem[ local_warp_id ]);

	int n = warp_id * (WARP_SIZE);
	bool finished = true;

	while (n < N) {
		int work_sz = ((N-n) > (WARP_SIZE)) ? (WARP_SIZE) : (N-n);

		// copy my work to shared memory 
		gpu_memcpy_SIMD2<int, int , WARP_SIZE> (local->levels, &(d_level[n]), work_sz, warp_offset);
		gpu_memcpy_SIMD<int, WARP_SIZE> (local->nodes, &(d_nodes[n]), work_sz+1, warp_offset);

		for(int j = 0; j < work_sz; j++)
		{

			if (local->levels[j] ==  level)  // found it
			{
				int from = local->nodes[j];
				int to = local->nodes[j + 1];

				// expansion is now done in diverge (SIMD)
				bool temp = expand_frontier_gpu_SIMD<WARP_SIZE>
					(M, from, to, level, d_edges, d_level, warp_offset, d_finished);
				finished = finished  && temp;

			}
		}

		n += grid_size;
  	}

    if (finished == false)
        *d_finished = false;
}
